# main.py

import argparse
import openai
from llama_index.llms.openai import OpenAI
from llama_index.core.response.notebook_utils import display_response
from llama_index.core import VectorStoreIndex

from src import config, data_handler, evaluator, finetuner


def setup():
    """初始化设置，如设置 API 密钥。"""
    if not config.OPENAI_API_KEY:
        raise ValueError("OPENAI_API_KEY 环境变量未设置。请在 .env 文件中设置。")
    openai.api_key = config.OPENAI_API_KEY
    print("OpenAI API 密钥已设置。")


def prepare_data_step():
    """执行所有数据准备步骤。"""
    print("\n--- 步骤 1: 准备数据 ---")
    data_handler.download_data()
    train_docs, eval_docs = data_handler.load_and_split_documents()

    # 为训练和评估生成问题
    data_handler.generate_and_save_questions(train_docs, config.TRAIN_QUESTIONS_PATH, config.NUM_QUESTIONS_TO_GENERATE)
    data_handler.generate_and_save_questions(eval_docs, config.EVAL_QUESTIONS_PATH, config.NUM_QUESTIONS_TO_GENERATE)
    print("数据准备完成。")


def evaluate_base_model_step():
    """评估基础模型的性能。"""
    print("\n--- 步骤 2: 评估基础模型 ---")
    _, all_documents = data_handler.load_and_split_documents()  # 使用所有文档构建索引
    eval_questions = data_handler.load_questions(config.EVAL_QUESTIONS_PATH)

    base_llm = OpenAI(model=config.BASE_MODEL, temperature=config.BASE_MODEL_TEMP)

    base_results = evaluator.evaluate_query_engine(all_documents, eval_questions, base_llm)

    print("\n基础模型评估结果:")
    print(base_results)


def finetune_step():
    """生成微调数据并执行微调。"""
    print("\n--- 步骤 3: 执行微调 ---")
    train_docs, _ = data_handler.load_and_split_documents()
    train_questions = data_handler.load_questions(config.TRAIN_QUESTIONS_PATH)

    # 1. 生成微调数据
    finetuner.generate_finetuning_data(train_docs, train_questions)

    # 2. 运行微调任务 (这将需要一些时间)
    # 注意：这一步会产生费用
    finetuner.run_finetuning_job()
    print("微调任务已完成。请记录下您的微调模型 ID。")


def evaluate_finetuned_model_step(model_id: str):
    """评估微调后模型的性能。"""
    if not model_id:
        print("错误：需要提供微调后的模型 ID。请使用 --model_id <your-ft-model-id>")
        return

    print(f"\n--- 步骤 4: 评估微调模型 (ID: {model_id}) ---")
    _, all_documents = data_handler.load_and_split_documents()
    eval_questions = data_handler.load_questions(config.EVAL_QUESTIONS_PATH)

    ft_llm = OpenAI(model=model_id, temperature=config.FINETUNED_MODEL_TEMP)

    ft_results = evaluator.evaluate_query_engine(all_documents, eval_questions, ft_llm)

    print("\n微调模型评估结果:")
    print(ft_results)


def compare_responses_step(model_id: str):
    """比较基础模型和微调模型的响应差异。"""
    if not model_id:
        print("错误：需要提供微调后的模型 ID。请使用 --model_id <your-ft-model-id>")
        return

    print("\n--- 步骤 5: 比较模型响应 ---")
    _, all_documents = data_handler.load_and_split_documents()
    eval_questions = data_handler.load_questions(config.EVAL_QUESTIONS_PATH)

    question_to_compare = eval_questions[12]
    print(f"用于比较的问题: '{question_to_compare}'")

    index = VectorStoreIndex.from_documents(all_documents)

    # 基础模型响应
    print("\n--- 原始 GPT-3.5-Turbo 响应 ---")
    base_llm = OpenAI(model=config.BASE_MODEL, temperature=config.BASE_MODEL_TEMP)
    query_engine_base = index.as_query_engine(llm=base_llm)
    response_base = query_engine_base.query(question_to_compare)
    display_response(response_base)

    # 微调模型响应
    print("\n--- 微调后模型响应 ---")
    ft_llm = OpenAI(model=model_id, temperature=config.FINETUNED_MODEL_TEMP)
    query_engine_ft = index.as_query_engine(llm=ft_llm)
    response_ft = query_engine_ft.query(question_to_compare)
    display_response(response_ft)


def main():
    parser = argparse.ArgumentParser(description="GPT-3.5-Turbo 微调项目工作流")
    parser.add_argument(
        "--step",
        type=str,
        required=True,
        choices=["prepare-data", "evaluate-base", "finetune", "evaluate-finetuned", "compare", "all"],
        help="要执行的工作流步骤"
    )
    parser.add_argument(
        "--model_id",
        type=str,
        help="微调后的模型ID（用于 evaluate-finetuned 和 compare 步骤）"
    )

    args = parser.parse_args()

    setup()

    if args.step == "prepare-data":
        prepare_data_step()
    elif args.step == "evaluate-base":
        evaluate_base_model_step()
    elif args.step == "finetune":
        finetune_step()
    elif args.step == "evaluate-finetuned":
        evaluate_finetuned_model_step(args.model_id)
    elif args.step == "compare":
        compare_responses_step(args.model_id)
    elif args.step == "all":
        # 注意: 'all' 模式需要手动交互，因为微调需要时间
        print("执行完整工作流...")
        prepare_data_step()
        evaluate_base_model_step()
        print("\n下一步是微调。这会产生费用并需要一些时间。")
        user_input = input("是否继续执行微调？ (y/n): ")
        if user_input.lower() == 'y':
            finetune_step()
            print("\n微调任务已完成。请从 OpenAI 平台获取您的模型 ID，")
            print("然后使用 'evaluate-finetuned' 和 'compare' 步骤来评估它。")
        else:
            print("已取消微调。")


if __name__ == "__main__":
    main()